import numpy as np
import torch
from torch.utils.data.distributed import DistributedSampler
from preprocess.energy import cal_energy_array


class Collator:
  def __init__(self, config):
    self.config = config
  def collate_rml(self, minibatch):
      config = self.config
      data = []
      label = []
      for record in minibatch:
          x = record['data']
          y = record['label']
          data.append(torch.tensor(x, dtype=torch.float32))
          label.append(torch.tensor(y, dtype=torch.long))

      data = torch.stack(data, dim=0)
      label = torch.stack(label, dim=0)
      return {'data': data, 'label': label.squeeze(), 'condition': None}

  def collate_rmlwithidx(self, minibatch):
      config = self.config
      data = []
      label = []
      index = []
      for record in minibatch:
          x = record['data']
          y = record['label']
          idx = record['index']
          data.append(torch.tensor(x, dtype=torch.float32))
          label.append(torch.tensor(y, dtype=torch.long))
          index.append(torch.tensor(idx, dtype=torch.long))

      data = torch.stack(data, dim=0)
      label = torch.stack(label, dim=0)
      index = torch.stack(index, dim=0)
      return {'data': data, 'label': label.squeeze(), 'index': index.squeeze()}


class RMLSNRDataset(torch.utils.data.Dataset):
  def __init__(self, data, le, config, stage='all'):
    super().__init__()
    self.config = config
    args = config.args
    if args.mode:
        if stage == 'specific':
            mask = data.df["SNR"] >= self.config.data.train.SNR
        elif stage == 'high':
            mask = data.df["SNR"] >= self.config.data.train.HSNR
        elif stage == 'medium':
            mask = (data.df["SNR"] >= self.config.data.train.LSNR) & (data.df["SNR"] < self.config.data.train.HSNR)
        elif stage == 'low':
            mask = data.df["SNR"] < self.config.data.train.LSNR
        else:
            mask = None
    else:
        if self.config.data.test.SNR is not None:
            mask = data.df["SNR"] == self.config.data.test.SNR
            # mask = (data.df["SNR"] == self.params.test_SNR) & (data.df["Modulation"] == '8PSK')
        else:
            mask = None

    self.x, self.y = data.as_numpy(le=le, mask=mask)

  def __len__(self):
    return len(self.x)

  def __getitem__(self, idx):
    x = self.x[idx]
    energy = cal_energy_array(x)
    x = x / np.sqrt(energy)
    y = self.y[idx]

    return {'data': x, 'label': y, 'index': idx}


class CollectedDataset(torch.utils.data.Dataset):
  def __init__(self, data_path, label_path, config):
    super().__init__()
    self.params = config
    self.x = np.load(data_path)  # Shape: [N, 1, 2, 128]
    self.y = np.load(label_path)  # Shape: [N,1]

  def __len__(self):
    return len(self.x)

  def __getitem__(self, idx):
    x = self.x[idx]
    energy = cal_energy_array(x)
    x = x / np.sqrt(energy)
    y = self.y[idx]

    return {'data': x, 'label': y, 'index': idx}


